#pragma once

#include <mcp/config.h>
#include <mcp/auth/authenticator.h>

#include <memory>
#include <string>
#include <vector>
#include <unordered_map>
#include <optional>
#include <functional>
#include <mutex>

namespace mcp {
namespace auth {

/// @brief Authentication middleware
class AuthHandler {
public:
    AuthHandler();
    ~AuthHandler();

    /// @brief Add an authenticator
    /// @param authenticator Authenticator to add
    void add_authenticator(std::unique_ptr<Authenticator> authenticator);

    /// @brief Remove authenticator by scheme
    /// @param scheme Authentication scheme
    void remove_authenticator(const std::string& scheme);

    /// @brief Authenticate request
    /// @param auth_header Authorization header
    /// @return Authentication result or error
    Result<AuthResult> authenticate(const std::string& auth_header);

    /// @brief Authenticate with credentials
    /// @param credentials Authentication credentials
    /// @return Authentication result or error
    Result<AuthResult> authenticate(const Credentials& credentials);

    /// @brief Generate authentication challenge
    /// @return WWW-Authenticate header value
    std::string generate_challenge() const;

    /// @brief Validate access token
    /// @param token Access token to validate
    /// @return true if valid, false otherwise
    bool validate_token(const std::string& token) const;

    /// @brief Set default authentication scheme
    /// @param scheme Default scheme
    void set_default_scheme(const std::string& scheme);

    /// @brief Get default authentication scheme
    /// @return Default scheme
    const std::string& get_default_scheme() const;

    /// @brief Get supported schemes
    /// @return List of supported authentication schemes
    std::vector<std::string> get_supported_schemes() const;

    /// @brief Check if authentication is required
    /// @return true if authentication is required, false otherwise
    bool is_authentication_required() const;

    /// @brief Set authentication requirement
    /// @param required Whether authentication is required
    void set_authentication_required(bool required);

    /// @brief Set authentication failure handler
    /// @param handler Function to handle authentication failures
    void set_failure_handler(std::function<void(const Error&)> handler);

private:
    struct Impl;
    std::unique_ptr<Impl> impl_;
};

/// @brief Token manager for handling access tokens
class TokenManager {
public:
    TokenManager();
    ~TokenManager();

    /// @brief Store access token
    /// @param token Access token
    /// @param expires_in Token lifetime in seconds
    void store_access_token(const std::string& token, int expires_in);

    /// @brief Store refresh token
    /// @param token Refresh token
    void store_refresh_token(const std::string& token);

    /// @brief Get access token
    /// @return Access token or empty string if not available
    std::string get_access_token() const;

    /// @brief Get refresh token
    /// @return Refresh token or empty string if not available
    std::string get_refresh_token() const;

    /// @brief Check if access token is valid
    /// @return true if valid, false otherwise
    bool is_access_token_valid() const;

    /// @brief Check if refresh token is available
    /// @return true if available, false otherwise
    bool has_refresh_token() const;

    /// @brief Clear all tokens
    void clear_tokens();

    /// @brief Load tokens from storage
    /// @param storage Storage backend
    /// @return true if loaded successfully, false otherwise
    bool load_tokens(std::function<std::string(const std::string&)> storage);

    /// @brief Save tokens to storage
    /// @param storage Storage backend
    /// @return true if saved successfully, false otherwise
    bool save_tokens(std::function<void(const std::string&, const std::string&)> storage) const;

    /// @brief Set token refresh callback
    /// @param callback Function to call when token needs refresh
    void set_refresh_callback(std::function<std::string(const std::string&)> callback);

    /// @brief Refresh access token
    /// @return true if refreshed successfully, false otherwise
    bool refresh_access_token();

private:
    struct Impl;
    std::unique_ptr<Impl> impl_;
};

/// @brief Authentication context for request processing
struct AuthContext {
    bool authenticated = false;
    std::string user_id;
    std::string username;
    std::vector<std::string> roles;
    std::unordered_map<std::string, std::string> claims;
    std::string access_token;
    std::chrono::system_clock::time_point expires_at;
};

/// @brief Authorization handler
class AuthorizationHandler {
public:
    AuthorizationHandler();
    ~AuthorizationHandler();

    /// @brief Check if user has permission
    /// @param context Authentication context
    /// @param permission Permission to check
    /// @return true if authorized, false otherwise
    bool has_permission(const AuthContext& context, const std::string& permission) const;

    /// @brief Check if user has role
    /// @param context Authentication context
    /// @param role Role to check
    /// @return true if user has role, false otherwise
    bool has_role(const AuthContext& context, const std::string& role) const;

    /// @brief Check if user has any of the required roles
    /// @param context Authentication context
    /// @param roles Roles to check
    /// @return true if user has any of the roles, false otherwise
    bool has_any_role(const AuthContext& context, const std::vector<std::string>& roles) const;

    /// @brief Check if user has all of the required roles
    /// @param context Authentication context
    /// @param roles Roles to check
    /// @return true if user has all of the roles, false otherwise
    bool has_all_roles(const AuthContext& context, const std::vector<std::string>& roles) const;

    /// @brief Add role-based access control rule
    /// @param role Role name
    /// @param permissions List of permissions for the role
    void add_role_permissions(const std::string& role, const std::vector<std::string>& permissions);

    /// @brief Add user-specific permissions
    /// @param user_id User ID
    /// @param permissions List of permissions for the user
    void add_user_permissions(const std::string& user_id, const std::vector<std::string>& permissions);

    /// @brief Clear all rules
    void clear_rules();

private:
    struct Impl;
    std::unique_ptr<Impl> impl_;
};

} // namespace auth
} // namespace mcp